# -*- encoding: utf-8 -*-

"""
@File:      hosts.py
@Time:      2022/05/12 16:09:15
@Author:    bolong.tbl
@Version:   1.0
@Contact:   bolong.tbl@alibaba-inc.com
@License:   Mulan PSL v2
"""

from avocado.utils import ssh
from avocado.utils import process


class Host(object):
    """
    Base Host class for all hosts
    """

    def __init__(self, host):
        self.host = host


class LocalHost(Host):
    """
    LocalHost class represents a localhost   

    Usage: 
        local = LocalHost()
    """

    def __init__(self, host='localhost'):
        super().__init__(host)

    def cmd(self, command, ignore_status=False):
        """
        Runs a command on localhost

        :param command (str): the command to execute on localhost
        :param ignore_status (bool): ignore exception if True, otherwise threw an exception. defaults to False.

        :returns (str): stdout if exit_status is 0, otherwise stderr
        """
        result = process.run(command, ignore_status=ignore_status, shell=True)

        if result.exit_status == 0:
            return result.exit_status, result.stdout_text.strip()
        else:
            return result.exit_status, result.stderr_text.strip()


class RemoteHost(Host):
    """
    RemoteHost class represents a remote host   

    Usage: 
        remote = RemoteHost(host='192.168.0.1',
                            port=22,
                            username='foo',
                            password='bar')

    """
    REBOOT_CMD = ['reboot', 'echoc>/proc/sysrq-trigger', 'kexec-e', 'halt--reboot', 'poweroff--reboot','reboot--force', 'shutdown-rnow', 'shutdown-r+1']

    def __init__(self, host, username, port=22, key=None, password=None):
        super().__init__(host)
        self.port = port
        self.username = username
        self.key = key
        self.password = password
        self.session = ssh.Session(host=self.host,
                                   port=self.port,
                                   user=self.username,
                                   key=self.key,
                                   password=self.password)
        self.session.connect()

    def cmd(self, command, ignore_status=False):
        """
        Runs a command over the ssh session

        Note: awk command with single quote ' is not supported

        :param command (str): the command to execute
        :param ignore_status (bool): ignore exception if True, otherwise threw an exception. defaults to False.

        :returns (str): stdout if exit_status is 0, otherwise stderr
        """
        if command.replace(' ', '') in self.REBOOT_CMD:
            result = process.run(self.session.get_raw_ssh_command(
                command), ignore_status=True, timeout=10)
        else:
            result = self.session.cmd(command, ignore_status=ignore_status)

        if result.exit_status == 0:
            return result.exit_status, result.stdout_text.strip()
        else:
            return result.exit_status, result.stderr_text.strip()


class VM(Host):
    pass


class Container(Host):
    pass
